{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Lab3_rl.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python2",
      "display_name": "Python 2"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "metadata": {
        "id": "WoXYKhfZMHiw",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "<table align=\"center\">\n",
        "  <td align=\"center\"><a target=\"_blank\" href=\"http://introtodeeplearning.com\">\n",
        "        <img src=\"http://introtodeeplearning.com/images/colab/mit.png\" style=\"padding-bottom:5px;\" />\n",
        "      Visit MIT Deep Learning</a></td>\n",
        "  <td align=\"center\"><a target=\"_blank\" href=\"https://colab.research.google.com/github/aamini/introtodeeplearning_labs/blob/master/lab3/Lab3_rl.ipynb\">\n",
        "        <img src=\"http://introtodeeplearning.com/images/colab/colab.png?v2.0\"  style=\"padding-bottom:5px;\" />Run in Google Colab</a></td>\n",
        "  <td align=\"center\"><a target=\"_blank\" href=\"https://github.com/aamini/introtodeeplearning_labs/blob/master/lab3/Lab3_rl.ipynb\">\n",
        "        <img src=\"http://introtodeeplearning.com/images/colab/github.png\"  height=\"70px\" style=\"padding-bottom:5px;\"  />View Source on GitHub</a></td>\n",
        "</table>\n",
        "\n",
        "# Lab 3: Model-Free Reinforcement Learning\n",
        "\n",
        "Reinforcement learning (RL) is a subset of machine learning which poses learning problems as interactions between agents and environments. It often assumes agents have no prior knowledge of the given world, so they must learn to navigate environments by optimizing some provided reward function. Within a world, an agent can take certain actions and receive feedback--in the form of positive or negative rewards--with respect to their decision. As such, an agent's feedback loop is somewhat akin to the manner in which a child might learn to distinguish between \"good\" and \"bad\" actions. In practical terms, our RL agent will interact with the environment by taking an action at each timestep, receiving a corresponding reward, and updating its state according to what it's \"learned\".  \n",
        "\n",
        "![alt text](https://www.kdnuggets.com/images/reinforcement-learning-fig1-700.jpg)\n",
        "\n",
        "## Why do we care about games? \n",
        "While the ultimate goal of reinforcement learning is to teach agents to act in the real, physical world, games provide a set of very useful properties that we also care about: \n",
        "\n",
        "1.   In many cases, games have perfectly describable enviornments. For example, all rules of chess can be formally written and programmed into a chess game simulator;\n",
        "2.   Massively parallelizable. Do not require running in the real world, therefore simultaneous environments can be run on large data clusters; \n",
        "3.   Fast prototyping of algorithms on simpler scenarios can speed up the development of algorithms that could eventually run in the real-world; and\n",
        "4.   ... Games are fun! \n",
        "\n",
        "In this lab, we focus on building a model-free reinforcement learning algorithm to master two different enviornments with varying complexity. \n",
        "\n",
        "1.   **Cartpole:   Balance a pole in an upright position by only moving your base left or right. Low-dimensional observation space.**\n",
        "2.   **Pong:   Beat a classical AI system designed at the game of Pong. High-dimensional observational space -- learning directly from raw pixels!  **\n"
      ]
    },
    {
      "metadata": {
        "id": "zmrHSiXKTXTY",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "#Part 1: Cartpole\n",
        "\n",
        "Since we're no longer dealing with a supervised learning task, let's take a step back and outline our approach to the problem. First, we'll need to create our environment and initialize our agent. Moreover, we'll need to provide our agent with some sort of mechanism for remembering action and reward history... in other words, a memory bank. Then we'll need to define our learning algorithm, much like we've done in previous labs.\n",
        "\n",
        "\n",
        "First we'll import TensorFlow, enable Eager execution, and also import some dependencies."
      ]
    },
    {
      "metadata": {
        "id": "xk5qeNPWCm00",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "!apt-get install -y xvfb python-opengl > /dev/null 2>&1\n",
        "!pip install gym pyvirtualdisplay scikit-video > /dev/null 2>&1\n",
        "\n",
        "import tensorflow as tf\n",
        "tf.enable_eager_execution()\n",
        "\n",
        "\n",
        "import gym\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from IPython import display as ipythondisplay\n",
        "import time\n",
        "\n",
        "# Download the class repository\n",
        "! git clone https://github.com/aamini/introtodeeplearning_labs.git  > /dev/null 2>&1\n",
        "% cd introtodeeplearning_labs \n",
        "! git pull\n",
        "% cd .. \n",
        "\n",
        "import introtodeeplearning_labs as util"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "UT7YL8KBJIIc",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### 1.1 Define and inspect the environment\n",
        "\n",
        "In order to model our environment, we'll be using a toolkit developed by OpenAI, [OpenAI Gym](https://gym.openai.com/). It provides several pre-defined environments for training and testing reinforcement learning agents, including those for classic physics control tasks, Atari video games, and robotic simulations. To access the basic version of a control task, \"Cart Pole\", we can use `env = gym.make(\"CartPole-v0\")`. When we imported `gym`, we gained access to higher level functions in the package, including creating virtual worlds. Each environment has a specific identifier (for which you can read through [here](https://gym.openai.com/envs/#classic_control)) which is accessed by passing the environment name as a string variable.\n",
        "One issue we might experience when developing RL algorithms is that many aspects of the learning process are inherently random: initializing game states, changes in the environment, and the agent's actions. As such, it can be helpful to set a random \"seed\" for one of these variables to ensure some level of reproducibility. Much like you might use `numpy.random.seed`, we can call the comparable function in gym, `seed`, with our defined environment to ensure the environment's random variables are initialized the same each time. "
      ]
    },
    {
      "metadata": {
        "id": "quv9SC0iIYFm",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "env = gym.make(\"CartPole-v0\")\n",
        "env.seed(1) # reproducible, since RL has high variance"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "mhEITUcKK455",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "\n",
        "**CartPole Environment: ** \n",
        "\n",
        "A pole is attached by an un-actuated joint to a cart, which moves along a frictionless track. The system is controlled by applying a force of +1 or -1 to the cart. The pendulum starts upright, and the goal is to prevent it from falling over. A reward of +1 is provided for every timestep that the pole remains upright. The episode ends when the pole is more than 15 degrees from vertical, or the cart moves more than 2.4 units from the center.\n",
        "\n",
        "<img width=\"400px\" src=\"https://danielpiedrahita.files.wordpress.com/2017/02/cart-pole.png\"></img>\n",
        "\n",
        "Observations:\n",
        "\n",
        "1. position of cart\n",
        "2. velocity of cart\n",
        "3. angle of pole\n",
        "4. rotation rate of pole\n",
        "\n",
        "We can confirm the size of the space by querying the observation space:\n"
      ]
    },
    {
      "metadata": {
        "id": "UVJaEcbdIX82",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "print \"Enviornment has observation space = {}\".format(env.observation_space)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "ZibGgjrALgPM",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "At every time step, the agent can move either right or left. Confirm the size of the action space by querying the environment:"
      ]
    },
    {
      "metadata": {
        "id": "qc9SIPxBIXrm",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "n_actions = env.action_space.n\n",
        "print \"Number of possible actions that the agent can choose from = {}\".format(n_actions)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "pPfHME8aRKkb",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### 1.2 Define the Agent\n",
        "\n",
        "Let's define our agent, which is simply a deep neural network which takes as input an observation of the enviornment and outputs the probability of taking each of the possible actions. \n"
      ]
    },
    {
      "metadata": {
        "id": "W-o_XK4oQ4eu",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def create_cartpole_model():\n",
        "  model = tf.keras.models.Sequential([\n",
        "      tf.keras.layers.Dense(units=32, activation='relu'),\n",
        "      # TODO: define the output dimension of the last Dense layer\n",
        "      # Hint: think about that the space the agent needs to act in\n",
        "      tf.keras.layers.Dense(units='''TODO''', activation=None) # TODO\n",
        "  ])\n",
        "  return model\n",
        "\n",
        "cartpole_model = create_cartpole_model()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "d5D5NSIYS2IW",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Define the action function that executes a forward pass through the network and samples from the output. Take special note of the output activation of the model."
      ]
    },
    {
      "metadata": {
        "id": "E_vVZRr8Q4R_",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def choose_action(model, observation):\n",
        "    \n",
        "  observation = observation.reshape([1, -1])\n",
        "  '''TODO: feed the observations through the model to predict the log probabilities of each possible action.'''\n",
        "  logits = # TODO\n",
        "  \n",
        "  # pass the log probabilities through a softmax to compute true probabilities\n",
        "  prob_weights = tf.nn.softmax(logits).numpy()\n",
        "  \n",
        "  '''TODO: randomly sample from the prob_weights to pick an action.\n",
        "  Hint: carefully consider the dimensionality of the input probabilities (vector) and the output action (scalar)'''\n",
        "  action = # TODO\n",
        "\n",
        "  return action"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "_tR9uAWcTnkr",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### 1.3 Create the agent's memory\n",
        "\n",
        "During training, the agent will need to remember all of its observations, actions so that once the episode ends, it can \"reinforce\" the good actions and punish the undesirable actions. Let's do this by defining a simple memory buffer that contains the agent's observations, actions, and received rewards from a given episode. "
      ]
    },
    {
      "metadata": {
        "id": "8MM6JwXVQ4JG",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "class Memory:\n",
        "  def __init__(self): \n",
        "      self.clear()\n",
        "\n",
        "  def clear(self): \n",
        "      self.observations = []\n",
        "      self.actions = []\n",
        "      self.rewards = []\n",
        "\n",
        "  def add_to_memory(self, new_observation, new_action, new_reward): \n",
        "      self.observations.append(new_observation)\n",
        "      '''TODO: update the list of actions with new action'''\n",
        "      # TODO\n",
        "      '''TODO: update the list of rewards with new reward'''\n",
        "      # TODO\n",
        "        \n",
        "memory = Memory()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "D4YhtPaUVj5m",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "We're almost ready to begin the learning algorithm for our agent! The final step is to compute the discounted rewards of our agent. Recall from lecture, we use reward discount to give more preference at getting rewards now rather than later in the future. The idea of discounting rewards is similar to discounting money in the case of interest and can be defined as: \n",
        "\n",
        ">$R_{t}=\\sum_{k=0}^\\infty\\gamma^kr_{t+k}$\n",
        "\n",
        "where  $\\gamma$ is the discount factor. In other words, at the end of an episode, we'll want to depreciate any rewards received at later time steps. Since we can't play an infinite number of games, we'll be limited to the number of timesteps in an episode. When implementing the function, you can initialize a numpy array of zeros (with length of the number of time steps) and fill it with the real discounted reward values as you loop through the saved rewards from the episode. We'll also want to normalize our output, which you can do using information about the mean and standard deviation of the discounted rewards.\n"
      ]
    },
    {
      "metadata": {
        "id": "5_Q2OFYtQ32X",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def normalize(x):\n",
        "  x -= np.mean(x)\n",
        "  x /= np.std(x)\n",
        "  return x\n",
        "\n",
        "def discount_rewards(rewards, gamma=0.95): \n",
        "  discounted_rewards = np.zeros_like(rewards)\n",
        "  R = 0\n",
        "  for t in reversed(range(0, len(rewards))):\n",
        "      # update the total discounted reward\n",
        "      R = R * gamma + rewards[t]\n",
        "      discounted_rewards[t] = R\n",
        "      \n",
        "  return normalize(discounted_rewards)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "QzbY-mjGYcmt",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### 1.4 Define the learning algorithm\n",
        "\n",
        "Now we can start to define the learing algorithm which will be used to reinforce good behaviors of the agent and discourage bad behaviours. Start by defining the optimizer we want to use."
      ]
    },
    {
      "metadata": {
        "id": "m3u6xDNMY0zg",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "learning_rate = 1e-3\n",
        "optimizer = tf.train.AdamOptimizer(learning_rate)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "M-LJwWqTZegG",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "And now let's define the loss function. In this lab we are focusing on policy gradient methods which aim to **maximize** the likelihood of actions that result in large rewards. Equivalently, this means that we want to **minimize** the negative likelihood of these same actions. Like in supervised learning, we can use stochastic gradient descent methods to achieve this minimization. \n",
        "\n",
        "Since the log function is monotonically increasing, this means that minimizing negative **likelihood** is equivalent to minimizing negative **log-likelihood**.  Recall that we can easily compute the negative log-likelihood of an discrete action by evaluting its softmax cross entropy (https://www.tensorflow.org/api_docs/python/tf/nn/sparse_softmax_cross_entropy_with_logits) "
      ]
    },
    {
      "metadata": {
        "id": "fsgZ3IDCY_Zn",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def compute_loss(logits, actions, rewards): \n",
        "  '''TODO: complete the call to compute the loss'''\n",
        "  neg_logprob = tf.nn.sparse_softmax_cross_entropy_with_logits('''TODO''', '''TODO''') # TODO\n",
        "  '''TODO: scale the negative log probability by the rewards'''\n",
        "  loss = tf.reduce_mean( '''TODO''' ) # TODO\n",
        "  return loss"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "Rr5vQ9fqbPpp",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Now let's use the loss function to define a backpropogation step of our learning algorithm."
      ]
    },
    {
      "metadata": {
        "id": "_50ada7nbZ7L",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def train_step(model, optimizer, observations, actions, discounted_rewards):\n",
        "  with tf.GradientTape() as tape:\n",
        "      # Forward propogate through the agent\n",
        "      observations = tf.convert_to_tensor(observations, dtype=tf.float32)\n",
        "      logits = model(observations)\n",
        "\n",
        "      '''TODO: call the compute_loss function to compute the loss'''\n",
        "      loss = # TODO\n",
        "\n",
        "  '''TODO: run backpropagation using the tape.gradient method'''\n",
        "  grads = tape.gradient('''TODO''', '''TODO''') # TODO\n",
        "  optimizer.apply_gradients(zip(grads, model.variables), global_step=tf.train.get_or_create_global_step())"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "XsjKXh6BcgjR",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### 1.5 Let the agent go and watch it learn from scratch!\n",
        "\n",
        "Having had no prior knowledge of the environment, the agent will begin to learn how to balance the pole on the cart based only on the feedback received from the environment! Having defined how our agent can move, how it takes in new observations, and how it updates its state, we'll see how it gradually learns a policy of actions to optimize balancing the pole as long as possible.\n",
        "\n"
      ]
    },
    {
      "metadata": {
        "id": "XmOzc2rrcn8Q",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "cartpole_model = create_cartpole_model()\n",
        "\n",
        "smoothed_reward = util.LossHistory(smoothing_factor=0.9)\n",
        "plotter = util.PeriodicPlotter(sec=5, xlabel='Iterations', ylabel='Rewards')\n",
        "\n",
        "\n",
        "for i_episode in range(1000):\n",
        "\n",
        "  plotter.plot(smoothed_reward.get())\n",
        "\n",
        "  # Restart the environment\n",
        "  observation = env.reset()\n",
        "\n",
        "  while True:\n",
        "      # using our observation, take an action\n",
        "      action = choose_action(cartpole_model, observation)\n",
        "      next_observation, reward, done, info = env.step(action)\n",
        "      # add to memory\n",
        "      memory.add_to_memory(observation, action, reward)\n",
        "      \n",
        "      # is the episode over? did you crash or do so well that you're done?\n",
        "      if done:\n",
        "          # determine total reward and keep a record of this\n",
        "          total_reward = sum(memory.rewards)\n",
        "          smoothed_reward.append( total_reward )\n",
        "          \n",
        "          # initiate training - remember we don't know anything about how the agent is doing until it's crashed!\n",
        "          train_step(cartpole_model, \n",
        "                     optimizer, \n",
        "                     observations = np.vstack(memory.observations),\n",
        "                     actions = np.array(memory.actions),\n",
        "                     discounted_rewards = discount_rewards(memory.rewards)\n",
        "                    )\n",
        "          \n",
        "          memory.clear()\n",
        "          break\n",
        "      # update our observatons\n",
        "      observation = next_observation"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "mkcUtGF1VE-K",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### 1.6 Save a video of the trained model while it is balancing the pole"
      ]
    },
    {
      "metadata": {
        "id": "M40RoTBxo3HD",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def save_video_of_model(model, env_name, filename='agent.mp4'):  \n",
        "  import skvideo.io\n",
        "  from pyvirtualdisplay import Display\n",
        "  display = Display(visible=0, size=(40, 30))\n",
        "  display.start()\n",
        "\n",
        "  env = gym.make(env_name)\n",
        "  obs = env.reset()\n",
        "  shape = env.render(mode='rgb_array').shape[0:2]\n",
        "\n",
        "  out = skvideo.io.FFmpegWriter(filename)\n",
        "\n",
        "  done = False\n",
        "  while not done: \n",
        "      frame = env.render(mode='rgb_array')\n",
        "      out.writeFrame(frame)\n",
        "      \n",
        "      action = model(tf.convert_to_tensor(obs.reshape((1,-1)), tf.float32)).numpy().argmax()\n",
        "      obs, reward, done, info = env.step(action)\n",
        "  out.close()\n",
        "  print \"Successfully saved into {}!\".format(filename)\n",
        "\n",
        "save_video_of_model(cartpole_model, \"CartPole-v0\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "dvvqdwO7VV_L",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### 1.7 Display the saved video\n"
      ]
    },
    {
      "metadata": {
        "id": "DBjhWQ0XwQ1d",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "from IPython.display import HTML\n",
        "import io, base64\n",
        "video = io.open('./agent.mp4', 'r+b').read()\n",
        "encoded = base64.b64encode(video)\n",
        "HTML(data='''\n",
        "<video controls>\n",
        "    <source src=\"data:video/mp4;base64,{0}\" type=\"video/mp4\" />\n",
        "</video>'''.format(encoded.decode('ascii')))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "CSbVNDpaVb3_",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Congratulations, well done! How does the agent perform? Could you train it for shorter amounts of time and still perform well? Would training longer help even more? "
      ]
    },
    {
      "metadata": {
        "id": "Eu6Mqxc720ST",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "#Part 2: Pong\n",
        "\n",
        "In Cart Pole, we dealt with an environment that was static--in other words, it didn't change over time. What happens if our environment is dynamic and unpredictable? Well that's exactly the case in Pong, since part of the environment is our opposing player. We don't know how our opponent will act or react to our actions, so the complexity of our problem increases. It also becomes much more interesting, since we can compete to beat our opponent."
      ]
    },
    {
      "metadata": {
        "id": "srZ4YE29isuA",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### 2.1 Define and inspect the environment"
      ]
    },
    {
      "metadata": {
        "id": "lbYHLr66i15n",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "env = gym.make(\"Pong-v0\")\n",
        "env.seed(1) # reproducible, since RL has high variance"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "52uZ2Xhyi-MW",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Observations: \n",
        "\n",
        "1. RGB image of shape (210, 160, 3)\n",
        "\n",
        "We can again confirm the size of the observation space by query:"
      ]
    },
    {
      "metadata": {
        "id": "0yX4GWvxjnHS",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "print \"Enviornment has observation space = {}\".format(env.observation_space)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "uuEC2TdSjx9D",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "At every time step, the agent has six actions to choose from: noop, fire, move right, move left, fire right, and fire left. Let's confirm the size of the action space by querying the environment:"
      ]
    },
    {
      "metadata": {
        "id": "Iuy9oPc1kag3",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "n_actions = env.action_space.n\n",
        "print \"Number of possible actions that the agent can choose from = {}\".format(n_actions)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "9-fghDRigUE5",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### 2.2 Define the Agent\n",
        "\n",
        "We'll define our agent again, but this time, we'll add convolutional layers to the network to increase the learning capacity of our network."
      ]
    },
    {
      "metadata": {
        "id": "IJiqbFYpgYRH",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def create_pong_model():\n",
        "  model = tf.keras.models.Sequential([\n",
        "      # Define and reshape inputs\n",
        "      tf.keras.layers.InputLayer(input_shape=(6400,), dtype=tf.float32),\n",
        "      tf.keras.layers.Reshape((80, 80, 1)),\n",
        "      \n",
        "      # Convolutional layers\n",
        "      tf.keras.layers.Conv2D(filters=16, kernel_size=(8,8), strides=(4,4), activation='relu', padding='same'),\n",
        "      # TODO: define a convolutional layer with 32 4x4 filters and 2x2 stride, ReLu activation\n",
        "      tf.keras.layers.Conv2D('''TODO''', '''TODO''', '''TODO''', '''TODO''', padding='same'), # TODO\n",
        "      tf.keras.layers.Flatten(),\n",
        "      \n",
        "      # Fully connected layer and output\n",
        "      tf.keras.layers.Dense(units=256, activation='relu'),\n",
        "      # TODO: define the output dimension of the last Dense layer\n",
        "      # Hint: think about that the space the agent needs to act in\n",
        "      tf.keras.layers.Dense('''TODO''', activation=None) # TODO\n",
        "  ])\n",
        "  return model\n",
        "\n",
        "pong_model = create_pong_model()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "yaeZ067olFiJ",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Since we've already defined the action function, `choose_action(model, observation)`, we don't need to define it again. Instead, we'll be able to reuse it later on by passing in our new model we've just created, `pong_model`. "
      ]
    },
    {
      "metadata": {
        "id": "l0RvqOVkmc2r",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### 2.3 Helper Functions"
      ]
    },
    {
      "metadata": {
        "id": "g4xtfog0mupM",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "We've already implemented some functions in Part 1 (Cartpole), so we won't need to recreate them in this section. However, we might need to make some slight modifications. For example, we need to reset the reward to zero when a game ends. In Pong, we know a game has ended if the reward is +1 (we won!) or -1 (we lost unfortunately). Otherwise, we expect the reward at a timestep to be zero. Also, note that we've increased gamma from 0.95 to 0.99, so the rate of decay will be even more rapid."
      ]
    },
    {
      "metadata": {
        "id": "iEZG2o50luLu",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def discount_rewards(rewards, gamma=0.99): \n",
        "  discounted_rewards = np.zeros_like(rewards)\n",
        "  R = 0\n",
        "  for t in reversed(range(0, len(rewards))):\n",
        "      # NEW: Reset sum\n",
        "      if rewards[t] != 0:\n",
        "        R = 0\n",
        "      # update the total discounted reward as before\n",
        "      R = R * gamma + rewards[t]\n",
        "      discounted_rewards[t] = R\n",
        "      \n",
        "  return normalize(discounted_rewards)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "HopLpb4IoOqA",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Before we input an image into our network, we'll need to pre-process it by converting it into a 1D array of floating point numbers:"
      ]
    },
    {
      "metadata": {
        "id": "Drpkn38Goout",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def pre_process(image):\n",
        "  I = image[35:195] # Crop\n",
        "  I = I[::2, ::2, 0] # Downsample width and height by a factor of 2\n",
        "  I[I == 144] = 0 # Remove background type 1\n",
        "  I[I == 109] = 0 # Remove background type 2\n",
        "  I[I != 0] = 1 # Set remaining elements (paddles, ball, etc.) to 1\n",
        "  return I.astype(np.float).ravel()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "-tP8_Bna6pgJ",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Let's use this function to visualize what an observation might look like before and after pre-processing"
      ]
    },
    {
      "metadata": {
        "id": "no5IIYtFm8pI",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "observation = env.reset()\n",
        "for i in range(30):\n",
        "  observation, _,_,_ = env.step(0)\n",
        "observation_pp = pre_process(observation)\n",
        "\n",
        "f = plt.figure(figsize=(10,3))\n",
        "ax = f.add_subplot(121)\n",
        "ax2 = f.add_subplot(122)\n",
        "ax.imshow(observation); ax.grid('off');\n",
        "ax2.imshow(observation_pp.reshape((80,80))); ax2.grid('off'); plt.title('Preprocessed Observation')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "mRqcaDQ1pm3x",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### 2.4: Training\n",
        "We've already defined our loss function with `compute_loss`, which is great! If we want to use a different learning rate, though, we can reinitialize the `optimizer`:"
      ]
    },
    {
      "metadata": {
        "id": "cIjRZ8JUqBLV",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "learning_rate=1e-4\n",
        "optimizer = tf.train.AdamOptimizer(learning_rate)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "IL2lMbTDn6Z3",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "We can also implement a very simple variant of `plot_progress`. In Pong, rather than feeding our network one image at a time, it can actually improve performance to input the difference between two consecutive observations, which really gives us information about the movement between frames. We'll first pre-process the raw observation, `x`, and then we'll compute the difference with the image frame we saw one timestep before. We'll also increase the number of maximum iterations from 1000 to 10000, since we expect it to take many more iterations to learn a more complex game."
      ]
    },
    {
      "metadata": {
        "id": "xCwyQQrPnkZG",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "pong_model = create_pong_model()\n",
        "MAX_ITERS = 10000\n",
        "\n",
        "smoothed_reward = util.LossHistory(smoothing_factor=0.9)\n",
        "plotter = util.PeriodicPlotter(sec=5, xlabel='Iterations', ylabel='Rewards')\n",
        "memory = Memory()\n",
        "\n",
        "for i_episode in range(MAX_ITERS):\n",
        "\n",
        "  plotter.plot(smoothed_reward.get())\n",
        "\n",
        "  # Restart the environment\n",
        "  observation = env.reset()\n",
        "  previous_frame = pre_process(observation)\n",
        "\n",
        "\n",
        "  while True:\n",
        "      # Pre-process image \n",
        "      current_frame = pre_process(observation)\n",
        "      \n",
        "      '''TODO: determine the observation change\n",
        "      Hint: this is the difference between the past two frames'''\n",
        "      obs_change = # TODO\n",
        "      \n",
        "      '''TODO: choose an action for the pong model, using the frame difference, and evaluate'''\n",
        "      action = # TODO \n",
        "      # Take the chosen action\n",
        "      next_observation, reward, done, info = env.step(action)\n",
        "      '''TODO: save the observed frame difference, the action that was taken, and the resulting reward!'''\n",
        "      # TODO\n",
        "      \n",
        "      # is the episode over? did you crash or do so well that you're done?\n",
        "      if done:\n",
        "          # determine total reward and keep a record of this\n",
        "          total_reward = sum(memory.rewards)\n",
        "          smoothed_reward.append( total_reward )\n",
        "          \n",
        "          # begin training\n",
        "          train_step(pong_model, \n",
        "                     optimizer, \n",
        "                     observations = np.vstack(memory.observations), \n",
        "                     actions = np.array(memory.actions),\n",
        "                     discounted_rewards = discount_rewards(memory.rewards))\n",
        "          \n",
        "          memory.clear()\n",
        "          break\n",
        "\n",
        "      observation = next_observation\n",
        "      previous_frame = current_frame"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "nwXjQH-puH5D",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### 2.5: Save and display video of training"
      ]
    },
    {
      "metadata": {
        "id": "8LiEY5Y_ts-Z",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "We can now save the video of our model learning:"
      ]
    },
    {
      "metadata": {
        "id": "TvHXbkL0tR6M",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "save_video_of_model(pong_model, \"Pong-v0\", filename='pong_agent.mp4')  "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "xmIcylIzuWaL",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "And display the result:"
      ]
    },
    {
      "metadata": {
        "id": "qoOBQSrXt2Ib",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "from IPython.display import HTML\n",
        "import io, base64\n",
        "video = io.open('./pong_agent.mp4', 'r+b').read()\n",
        "encoded = base64.b64encode(video)\n",
        "HTML(data='''\n",
        "<video controls>\n",
        "    <source src=\"data:video/mp4;base64,{0}\" type=\"video/mp4\" />\n",
        "</video>'''.format(encoded.decode('ascii')))"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}